"""
期货经纪人模块
"""

from .broker import Broker, OrderType, OrderAction
from loguru import logger


class FuturesBroker(Broker):
    """
    期货经纪人类，继承自Broker类，增加了期货特有的功能
    """
    
    def __init__(self, initial_cash=100000.0, commission=0.0, timeline=None, ticker_info_file=None):
        """
        初始化期货经纪人
        
        Parameters
        ----------
        initial_cash : float, optional
            初始现金，默认为100000.0
        commission : float, optional
            佣金比例，默认为0.0
        timeline : pandas.Series, optional
            时间序列，包含所有可能的交易时间点
        ticker_info_file : str, optional
            交易品种信息文件路径，如果提供则自动加载品种信息
        """
        super().__init__(initial_cash, commission, timeline, ticker_info_file)
        
        # 保证金设置
        self._margin_rates = {}  # instrument -> 保证金比例的映射
        # 合约乘数设置
        self._multipliers = {}   # instrument -> 合约乘数的映射
        
        # 如果有交易品种信息，自动设置保证金比例和合约乘数
        if self._ticker_info is not None:
            for inst, info in self._ticker_info.iterrows():
                self.set_margin_rate(inst, info['margin_rate'])
                self.set_multiplier(inst, info['multiplier'])
    
    def set_margin_rate(self, instrument, rate):
        """
        设置品种的保证金比例
        
        Parameters
        ----------
        instrument : str
            交易品种代码
        rate : float
            保证金比例
        """
        self._margin_rates[instrument] = rate
    
    def get_margin_rate(self, instrument):
        """
        获取品种的保证金比例
        
        Parameters
        ----------
        instrument : str
            交易品种代码
            
        Returns
        -------
        float
            保证金比例，如果未设置则返回0.1
        """
        return self._margin_rates.get(instrument, 0.1)
    
    def set_multiplier(self, instrument, multiplier):
        """
        设置品种的合约乘数
        
        Parameters
        ----------
        instrument : str
            交易品种代码
        multiplier : float
            合约乘数
        """
        self._multipliers[instrument] = multiplier
    
    def get_multiplier(self, instrument):
        """
        获取品种的合约乘数
        
        Parameters
        ----------
        instrument : str
            交易品种代码
            
        Returns
        -------
        float
            合约乘数，如果未设置则返回1.0
        """
        return self._multipliers.get(instrument, 1.0)
    
    def _calculate_margin(self, instrument, quantity, price):
        """
        计算保证金
        
        Parameters
        ----------
        instrument : str
            交易品种代码
        quantity : float
            数量
        price : float
            价格
            
        Returns
        -------
        float
            保证金金额
        """
        return (
            abs(quantity) * 
            price * 
            self.get_multiplier(instrument) * 
            self.get_margin_rate(instrument)
        )
    
    def order(self, instrument, quantity, order_type=OrderType.MARKET, price=None):
        """
        下单
        
        Parameters
        ----------
        instrument : str
            交易品种代码
        quantity : float
            数量，正数表示买入，负数表示卖出
        order_type : OrderType, optional
            订单类型，默认为市价单
        price : float, optional
            价格，对于限价单是必须的
            
        Returns
        -------
        int
            订单ID，如果资金不足返回None
        """
        # 检查资金是否足够
        estimated_price = price if price is not None else self._get_last_price(instrument)
        if estimated_price is None:
            logger.error(f"无法获取{instrument}的价格信息，订单创建失败")
            return None
            
        # 计算保证金和手续费
        margin = self._calculate_margin(instrument, quantity, estimated_price)
        commission = (
            abs(quantity) * 
            estimated_price * 
            self.get_multiplier(instrument) * 
            self._commission
        )
        
        # 检查可用资金
        required_capital = margin + commission
        # TODO 这个判断只是判断了一个，如果一个bar多个合约开仓，就可能有问题。需要在订单提交的时候预扣减，然后处理完了之后再把预扣减变为真实扣减
        if required_capital > self._cash:
            logger.warning(
                f"资金不足，订单创建失败。所需资金: {required_capital:.2f}, "
                f"可用资金: {self._cash:.2f}"
            )
            return None
        
        return super().order(instrument, quantity, order_type, price)
    
    def _update_position(self, order, fill_price, commission):
        """
        更新持仓
        
        Parameters
        ----------
        order : Order
            订单对象
        fill_price : float
            成交价格
        commission : float
            佣金
        """
        instrument = order.instrument
        quantity = order.quantity
        if order.action == OrderAction.SELL:
            quantity = -quantity
        
        multiplier = self.get_multiplier(instrument)
        margin = self._calculate_margin(instrument, quantity, fill_price)
        
        old_position = (
            self._positions.loc[instrument] 
            if instrument in self._positions.index 
            else None
        )
        
        if old_position is None or old_position['quantity'] == 0:
            # 新开仓
            self._positions.loc[instrument] = {
                'quantity': quantity,
                'avg_price': fill_price,
                'cost_basis': margin,
                'realized_pnl': 0,
                'unrealized_pnl': 0
            }
            self._cash -= margin + commission
            logger.info(
                f"开仓: {instrument}, 数量={abs(quantity)}, "
                f"价格={fill_price:.2f}, 保证金={margin:.2f}"
            )
            return
        
        old_quantity = old_position['quantity']
        old_avg_price = old_position['avg_price']
        new_quantity = old_quantity + quantity
        
        if new_quantity == 0:
            # 完全平仓
            realized_pnl = (
                (fill_price - old_avg_price) * 
                old_quantity * 
                multiplier
            ) - commission
            # 记录平仓信息
            logger.info(
                f"平仓: {instrument}, 数量={abs(old_quantity)}, "
                f"成交价={fill_price:.2f}, 盈亏={realized_pnl:.2f}"
            )
            # 返还保证金
            returned_margin = abs(old_quantity) * old_avg_price * multiplier * self.get_margin_rate(instrument)
            self._cash += returned_margin
            # 删除持仓记录
            self._positions.drop(instrument, inplace=True)
        elif (old_quantity > 0 and quantity > 0) or (old_quantity < 0 and quantity < 0):
            # 同向加仓
            new_avg_price = (
                old_quantity * old_avg_price + quantity * fill_price
            ) / new_quantity
            self._positions.loc[instrument, 'quantity'] = new_quantity
            self._positions.loc[instrument, 'avg_price'] = new_avg_price
            self._positions.loc[instrument, 'cost_basis'] += margin
            self._cash -= margin + commission
            logger.info(
                f"加仓: {instrument}, 数量={abs(quantity)}, "
                f"均价={new_avg_price:.2f}, 总持仓={new_quantity}"
            )
        else:
            # 反向开仓或部分平仓
            if abs(old_quantity) > abs(quantity):
                # 部分平仓
                realized_pnl = (
                    (fill_price - old_avg_price) * 
                    quantity * 
                    multiplier
                ) - commission
                # 返还部分保证金
                returned_margin = abs(quantity) * old_avg_price * multiplier * self.get_margin_rate(instrument)
                self._cash += returned_margin
                
                self._positions.loc[instrument, 'quantity'] = new_quantity
                self._positions.loc[instrument, 'cost_basis'] = (
                    abs(new_quantity) * 
                    old_avg_price * 
                    multiplier * 
                    self.get_margin_rate(instrument)
                )
                self._positions.loc[instrument, 'realized_pnl'] += realized_pnl
                
                logger.info(
                    f"部分平仓: {instrument}, 数量={abs(quantity)}, "
                    f"成交价={fill_price:.2f}, 盈亏={realized_pnl:.2f}, "
                    f"剩余持仓={new_quantity}"
                )
            else:
                # 完全平仓后反向开仓
                realized_pnl = (
                    (fill_price - old_avg_price) * 
                    old_quantity * 
                    multiplier
                ) - commission
                # 返还旧保证金并扣除新保证金
                returned_margin = abs(old_quantity) * old_avg_price * multiplier * self.get_margin_rate(instrument)
                self._cash += returned_margin - margin
                
                self._positions.loc[instrument, 'quantity'] = new_quantity
                self._positions.loc[instrument, 'avg_price'] = fill_price
                self._positions.loc[instrument, 'cost_basis'] = margin
                self._positions.loc[instrument, 'realized_pnl'] += realized_pnl
                
                logger.info(
                    f"平仓并反向开仓: {instrument}, 平仓数量={abs(old_quantity)}, "
                    f"开仓数量={abs(new_quantity)}, 成交价={fill_price:.2f}, "
                    f"平仓盈亏={realized_pnl:.2f}"
                ) 